package net.cyue.ort.llm.data;

import java.util.Collections;
import java.util.Map;

/**
 * 模型输出数据类
 * 封装模型推理的输出结果，包括logits和更新的缓存
 */
public class ModelOutput {
    private final float[][][] logits;
    private final Map<String, Object> pastKeyValueCache;
    
    public ModelOutput(float[][][] logits, Map<String, Object> pastKeyValueCache) {
        this.logits = logits;
        this.pastKeyValueCache = pastKeyValueCache != null ? pastKeyValueCache : Collections.emptyMap();
    }
    
    public float[][][] getLogits() {
        return logits;
    }
    
    public Map<String, Object> getPastKeyValueCache() {
        return pastKeyValueCache;
    }
    
    /**
     * 获取最后一个token位置的logits
     * @return 最后一个token的logits数组
     */
    public float[] getLastTokenLogits() {
        if (logits == null || logits.length == 0 || 
            logits[0] == null || logits[0].length == 0) {
            return new float[0];
        }
        return logits[0][logits[0].length - 1];
    }
    
    @Override
    public String toString() {
        return "ModelOutput{" +
            "logitsShape=" + (logits != null ? 
                "[" + logits.length + "][" + 
                (logits.length > 0 && logits[0] != null ? logits[0].length : 0) + "][" +
                (logits.length > 0 && logits[0] != null && logits[0].length > 0 && logits[0][0] != null ? 
                    logits[0][0].length : 0) + "]" : "null") +
            ", hasCache=" + !pastKeyValueCache.isEmpty() +
            '}';
    }
}

